from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from patagent.llm import PatsnapLLM
from patagent.prompt import MATERIAL_TRANSLATION_FEWSHOT_PROMPT


llm = PatsnapLLM()

examples = [
    {"input": "Barium bismuth strontium titanium oxide-Ba<sub>0.48</sub>Bi<sub>0.4</sub>Sr<sub>0.32</sub>Ti<sub>1.4</sub>O<sub>4.2</sub>", "output": "钡铋锶钛氧化物-Ba<sub>0.48</sub>Bi<sub>0.4</sub>Sr<sub>0.32</sub>Ti<sub>1.4</sub>O<sub>4.2</sub>"},
    {"input": "Abernathyite-K[U(AsO<sub>4</sub>)O<sub>2</sub>].4H<sub>2</sub>O", "output": "阿伯纳西石-K[U(AsO<sub>4</sub>)O<sub>2</sub>].4H<sub>2</sub>O"},
    {"input": "Helium; potassium-He<sub>2</sub><sup>39</sup>K", "output": "氦；钾-He<sub>2</sub><sup>39</sup>K"},
    {"input": "Acmite, calcian magnesian manganoan", "output": "钙质镁质锰质霓辉石"},
]

# Prompt
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}"),
    ]
)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)
translation_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", MATERIAL_TRANSLATION_FEWSHOT_PROMPT),
        few_shot_prompt,
        ("human", "question: \n\n {question}"),
    ]
)

# Chain
translation_chain = translation_prompt | llm | StrOutputParser()